import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse
import logging
from pathlib import Path
from image_style_trans.style_processor import StyleProcessor

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

def process_single_image(args):
    """
    处理单个图像
    """
    # 初始化风格处理器
    processor = StyleProcessor(api_key=args.api_key)
    
    if not processor.available:
        logging.error("风格处理器初始化失败，请检查API密钥和OSS配置")
        return False
    
    # 处理图像
    result_path = processor.process_image(
        image_path=args.image,
        prompt=args.prompt,
        negative_prompt=args.negative,
        output_dir=args.output_dir
    )
    
    if result_path:
        logging.info(f"图像处理成功，结果保存在: {result_path}")
        return True
    else:
        logging.error("图像处理失败")
        return False

def process_directory(args):
    """
    处理目录中的所有图像
    """
    # 初始化风格处理器
    processor = StyleProcessor(api_key=args.api_key)
    
    if not processor.available:
        logging.error("风格处理器初始化失败，请检查API密钥和OSS配置")
        return False
    
    # 获取目录中的所有图像
    input_dir = Path(args.directory)
    if not input_dir.exists() or not input_dir.is_dir():
        logging.error(f"输入目录不存在或不是目录: {input_dir}")
        return False
    
    # 根据文件扩展名筛选图像文件
    image_extensions = [".jpg", ".jpeg", ".png"]
    image_paths = []
    for ext in image_extensions:
        image_paths.extend(list(input_dir.glob(f"*{ext}")))
        image_paths.extend(list(input_dir.glob(f"*{ext.upper()}")))
    
    if not image_paths:
        logging.error(f"目录中未找到图像文件: {input_dir}")
        return False
    
    logging.info(f"找到{len(image_paths)}个图像文件")
    
    # 创建输出目录
    output_dir = args.output_dir or input_dir / "styled"
    os.makedirs(output_dir, exist_ok=True)
    
    # 批量处理图像
    successful_paths = processor.process_batch(
        image_paths=[str(p) for p in image_paths],
        prompt=args.prompt,
        negative_prompt=args.negative,
        output_dir=output_dir
    )
    
    if successful_paths:
        logging.info(f"批量处理完成，成功处理{len(successful_paths)}/{len(image_paths)}个图像")
        logging.info(f"处理结果保存在: {output_dir}")
        return True
    else:
        logging.error("批量处理失败，未成功处理任何图像")
        return False

def main():
    parser = argparse.ArgumentParser(description="图像风格转换工具")
    subparsers = parser.add_subparsers(dest="command", help="子命令")
    
    # 单个图像处理子命令
    single_parser = subparsers.add_parser("single", help="处理单个图像")
    single_parser.add_argument("--image", "-i", required=True, help="输入图像路径")
    single_parser.add_argument("--prompt", "-p", required=True, help="风格描述提示词")
    single_parser.add_argument("--negative", "-n", help="负向提示词")
    single_parser.add_argument("--output-dir", "-o", help="输出目录")
    single_parser.add_argument("--api-key", "-k", help="可灵AI API密钥")
    
    # 目录批处理子命令
    dir_parser = subparsers.add_parser("directory", help="处理目录中的所有图像")
    dir_parser.add_argument("--directory", "-d", required=True, help="输入目录路径")
    dir_parser.add_argument("--prompt", "-p", required=True, help="风格描述提示词")
    dir_parser.add_argument("--negative", "-n", help="负向提示词")
    dir_parser.add_argument("--output-dir", "-o", help="输出目录")
    dir_parser.add_argument("--api-key", "-k", help="可灵AI API密钥")
    
    args = parser.parse_args()
    
    if args.command == "single":
        success = process_single_image(args)
    elif args.command == "directory":
        success = process_directory(args)
    else:
        parser.print_help()
        return
    
    # 设置退出码
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()